import base64
from flask import request, jsonify, Blueprint
import cv2
import tempfile
from server.result.result import Result
from ultralytics import YOLOv10
import os

yolo_api = Blueprint('yolo_api', __name__)


@yolo_api.route("/yolo", methods=['POST'])
def yolov10_inference():
    # model = YOLOv10.from_pretrained(f'pt/{model_id}')
    data = request.get_json()
    image = data['image']
    video = data['video']
    # model_id = data['model_id']
    image_size = data['image_size']
    conf_threshold = data['conf_threshold']

    model = YOLOv10(r'/pt/yolov10n.pt')
    if image:
        results = model.predict(source=image, imgsz=image_size, conf=conf_threshold, save=True)
        print(results)
        annotated_image = results[0].plot()
        retval, buffer = cv2.imencode('.webp', annotated_image)
        image_bytes = buffer.tobytes()
        # os.remove(image)
        base64_image = base64.b64encode(image_bytes).decode('utf-8')
        response_data = {
            'image': f'data:image/webp;base64,{base64_image}'
        }
        # return send_file(BytesIO(image_bytes), mimetype='image/webp')
        # return annotated_image[:, :, ::-1], None
        # return annotated_image, None
        return jsonify(response_data)

    else:
        video_path = tempfile.mktemp(suffix=".webm")
        with open(video_path, "wb") as f:
            with open(video, "rb") as g:
                f.write(g.read())

        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        output_video_path = tempfile.mktemp(suffix=".webm")
        out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'vp80'), fps, (frame_width, frame_height))

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold)
            annotated_frame = results[0].plot()
            out.write(annotated_frame)

        cap.release()
        out.release()
        os.remove(video)

        return None, output_video_path


@yolo_api.route("/upload", methods=['POST'])
def upload_image():
    # 生成保存文件的目录，如果不存在则创建
    save_dir = '../../uploads'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    # 清除save_dir下的文件
    for filename in os.listdir(save_dir):
        file_path = os.path.join(save_dir, filename)
        try:
            # 删除文件
            os.remove(file_path)
        except Exception as e:
            print(f"删除文件 {file_path} 时出错: {e}")

    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    if file:
        # 保存文件
        filename = file.filename
        file_path = os.path.join(save_dir, filename)
        file.save(file_path)
        return jsonify({'message': 'File uploaded successfully', 'file_path': file_path}), 200


@yolo_api.route("/api_test", methods=["GET", "POST"])
def my_test():
    if request.method == "GET":
        print("访问成功")
        return Result.success()
    elif request.method == "POST":
        data = request.get_json()
        print(data)
        return Result.success_data(data)
